// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/util/arbitrator_test_util.h>
#include <pollux/common/memory/shared_arbitrator.h>
#include <pollux/dwio/dwrf/common/config.h>
#include <pollux/exec/table_writer.h>
#include <turbo/log/logging.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::exec::test;
using namespace kumo::pollux::memory;
using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    std::shared_ptr<core::QueryCtx> newQueryCtx(
        MemoryManager *memoryManager,
        melon::Executor *executor,
        int64_t memoryCapacity,
        const std::string &queryId) {
        std::unordered_map<std::string, std::shared_ptr<config::ConfigBase> > configs;
        std::shared_ptr<MemoryPool> pool =
                memoryManager->addRootPool("", memoryCapacity);
        auto queryCtx = core::QueryCtx::create(
            executor,
            core::QueryConfig({}),
            configs,
            cache::AsyncDataCache::getInstance(),
            std::move(pool),
            nullptr,
            queryId);
        return queryCtx;
    }

    std::unique_ptr<memory::MemoryManager> createMemoryManager(
        int64_t arbitratorCapacity,
        uint64_t memoryPoolInitCapacity,
        uint64_t maxArbitrationTimeMs,
        uint64_t fastExponentialGrowthCapacityLimit,
        double slowCapacityGrowPct) {
        memory::MemoryManagerOptions options;
        options.arbitratorCapacity = arbitratorCapacity;
        // Avoid allocation failure in unit tests.
        options.allocatorCapacity = arbitratorCapacity * 2;
        options.arbitratorKind = "SHARED";
        options.checkUsageLeak = true;
        using ExtraConfig = SharedArbitrator::ExtraConfig;
        options.extraArbitratorConfigs = {
            {
                std::string(ExtraConfig::kMemoryPoolInitialCapacity),
                melon::to<std::string>(memoryPoolInitCapacity) + "B"
            },
            {
                std::string(ExtraConfig::kMaxMemoryArbitrationTime),
                melon::to<std::string>(maxArbitrationTimeMs) + "ms"
            },
            {std::string(ExtraConfig::kGlobalArbitrationEnabled), "true"},
            {
                std::string(ExtraConfig::kFastExponentialGrowthCapacityLimit),
                melon::to<std::string>(fastExponentialGrowthCapacityLimit) + "B"
            },
            {
                std::string(ExtraConfig::kSlowCapacityGrowPct),
                melon::to<std::string>(slowCapacityGrowPct)
            }
        };
        options.arbitrationStateCheckCb = memoryArbitrationStateCheck;
        return std::make_unique<memory::MemoryManager>(options);
    }

    core::PlanNodePtr hashJoinPlan(
        const std::vector<RowVectorPtr> &vectors,
        core::PlanNodeId &joinNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors, true)
                .project({"c0", "c1", "c2"})
                .hashJoin(
                    {"c0"},
                    {"u1"},
                    PlanBuilder(planNodeIdGenerator)
                    .values(vectors, true)
                    .project({"c0 AS u0", "c1 AS u1", "c2 AS u2"})
                    .planNode(),
                    "",
                    {"c0", "c1", "c2"},
                    core::JoinType::kInner)
                .capturePlanNodeId(joinNodeId)
                .planNode();
    }

    QueryTestResult runHashJoinTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        bool enableSpilling,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto plan = hashJoinPlan(vectors, result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, true)
                    .config(core::QueryConfig::kJoinSpillEnabled, true)
                    .config(core::QueryConfig::kSpillStartPartitionBit, "29")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    core::PlanNodePtr aggregationPlan(
        const std::vector<RowVectorPtr> &vectors,
        core::PlanNodeId &aggregateNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors)
                .singleAggregation({"c0", "c1"}, {"array_agg(c2)"})
                .capturePlanNodeId(aggregateNodeId)
                .planNode();
    }

    QueryTestResult runAggregateTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        bool enableSpilling,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto plan = aggregationPlan(vectors, result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data =
                    AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, "true")
                    .config(core::QueryConfig::kAggregationSpillEnabled, "true")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    core::PlanNodePtr orderByPlan(
        const std::vector<RowVectorPtr> &vectors,
        core::PlanNodeId &orderNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors)
                .project({"c0", "c1", "c2"})
                .orderBy({"c2 ASC NULLS LAST"}, false)
                .capturePlanNodeId(orderNodeId)
                .planNode();
    }

    QueryTestResult runOrderByTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        bool enableSpilling,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto plan = orderByPlan(vectors, result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, "true")
                    .config(core::QueryConfig::kOrderBySpillEnabled, "true")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    core::PlanNodePtr rowNumberPlan(
        const std::vector<RowVectorPtr> &vectors,
        core::PlanNodeId &rowNumberNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors)
                .rowNumber({"c0"}, 2, false)
                .project({"c0", "c1"})
                .capturePlanNodeId(rowNumberNodeId)
                .planNode();
    }

    QueryTestResult runRowNumberTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        bool enableSpilling,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto plan = rowNumberPlan(vectors, result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, "true")
                    .config(core::QueryConfig::kRowNumberSpillEnabled, "true")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    core::PlanNodePtr topNPlan(
        const std::vector<RowVectorPtr> &vectors,
        core::PlanNodeId &topNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors)
                .project({"c1"})
                .topN({"c1 NULLS FIRST"}, 10, false)
                .capturePlanNodeId(topNodeId)
                .planNode();
    }

    QueryTestResult runTopNTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        bool enableSpilling,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto plan = topNPlan(vectors, result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data =
                    AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, "true")
                    .config(core::QueryConfig::kTopNRowNumberSpillEnabled, "true")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    core::PlanNodePtr writePlan(
        const std::vector<RowVectorPtr> &vectors,
        const std::string &outputDirPath,
        core::PlanNodeId &writeNodeId) {
        auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
        return PlanBuilder(planNodeIdGenerator)
                .values(vectors)
                .tableWrite(outputDirPath)
                .singleAggregation(
                    {}, {fmt::format("sum({})", TableWriteTraits::rowCountColumnName())})
                .capturePlanNodeId(writeNodeId)
                .planNode();
    }

    QueryTestResult runWriteTask(
        const std::vector<RowVectorPtr> &vectors,
        const std::shared_ptr<core::QueryCtx> &queryCtx,
        bool serialExecution,
        uint32_t numDrivers,
        memory::MemoryPool *pool,
        const std::string &kHiveConnectorId,
        bool enableSpilling,
        const RowVectorPtr &expectedResult) {
        QueryTestResult result;
        const auto outputDirectory = TempDirectoryPath::create();
        auto plan = writePlan(vectors, outputDirectory->getPath(), result.planNodeId);
        if (enableSpilling) {
            const auto spillDirectory = exec::test::TempDirectoryPath::create();
            result.data =
                    AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .spillDirectory(spillDirectory->getPath())
                    .config(core::QueryConfig::kSpillEnabled, "true")
                    .config(core::QueryConfig::kAggregationSpillEnabled, "false")
                    .config(core::QueryConfig::kWriterSpillEnabled, "true")
                    // Set 0 file writer flush threshold to always trigger flush in
                    // test.
                    .config(core::QueryConfig::kWriterFlushThresholdBytes, "0")
                    // Set stripe size to extreme large to avoid writer internal
                    // triggered flush.
                    .connectorSessionProperty(
                        kHiveConnectorId,
                        dwrf::Config::kOrcWriterMaxStripeSizeSession,
                        "1GB")
                    .connectorSessionProperty(
                        kHiveConnectorId,
                        dwrf::Config::kOrcWriterMaxDictionaryMemorySession,
                        "1GB")
                    .connectorSessionProperty(
                        kHiveConnectorId,
                        dwrf::Config::kOrcWriterMaxDictionaryMemorySession,
                        "1GB")
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        } else {
            result.data = AssertQueryBuilder(plan)
                    .serialExecution(serialExecution)
                    .queryCtx(queryCtx)
                    .maxDrivers(numDrivers)
                    .copyResults(pool, result.task);
        }
        if (expectedResult != nullptr) {
            assertEqualResults({result.data}, {expectedResult});
        }
        return result;
    }

    TestSuspendedSection::TestSuspendedSection(Driver *driver) : driver_(driver) {
        if (driver->task()->enterSuspended(driver->state()) != StopReason::kNone) {
            POLLUX_FAIL("Terminate detected when entering suspended section");
        }
    }

    TestSuspendedSection::~TestSuspendedSection() {
        if (driver_->task()->leaveSuspended(driver_->state()) != StopReason::kNone) {
            KLOG(WARNING)
                    << "Terminate detected when leaving suspended section for driver "
                    << driver_->driverCtx()->driverId << " from task "
                    << driver_->task()->taskId();
        }
    }
} // namespace kumo::pollux::exec::test
